{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4aea3368",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "e:\\Anaconda\\Anaconda\\envs\\py310\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import spacy\n",
    "import chromadb\n",
    "import re\n",
    "\n",
    "from sentence_transformers import SentenceTransformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e035fcc",
   "metadata": {},
   "source": [
    "## spyCa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3ba36e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"zh_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "60123864",
   "metadata": {},
   "outputs": [],
   "source": [
    "def semantic_chunking_with_sliding_window(\n",
    "    text: str,\n",
    "    chunk_size: int = 100,\n",
    "    overlap: int = 10\n",
    ") -> list:\n",
    "    \"\"\" \n",
    "    基于句子边界进行语义分块 + 滑动窗口策略（使用句子作为overlap单位）\n",
    "    \n",
    "    参数:\n",
    "        text (str): 输入文本\n",
    "        chunk_size (int): 每个chunk的最大字符数\n",
    "        overlap (int): 相邻chunk之间的字符级重叠区域（通过句子近似控制）\n",
    "    \n",
    "    返回:\n",
    "        List[str]: 按滑动窗口切分后的文本块列表\n",
    "    \"\"\"\n",
    "    # Step 1: 文本预处理\n",
    "    text = re.sub(r'\\s+', ' ', text).strip()\n",
    "\n",
    "    # Step 2: 使用 spaCy 分句\n",
    "    doc = nlp(text)\n",
    "    sentences = [sent.text.strip() for sent in doc.sents if sent.text.strip()]\n",
    "    # print(len(sentences))\n",
    "\n",
    "    # Step 3: 滑动窗口切分（按句子长度堆叠）\n",
    "    chunks = []\n",
    "    i = 0\n",
    "    while i < len(sentences):\n",
    "        current_chunk = []\n",
    "        current_length = 0\n",
    "        start_i = i\n",
    "        while i < len(sentences) and current_length + len(sentences[i]) <= chunk_size:\n",
    "            current_chunk.append(sentences[i])\n",
    "            current_length += len(sentences[i])\n",
    "            i += 1\n",
    "        chunks.append(' '.join(current_chunk))\n",
    "        # 回退以实现 overlap（按句子单位回退）\n",
    "        overlap_len = 0\n",
    "        j = len(current_chunk) - 1\n",
    "        while j >= 0 and overlap_len + len(current_chunk[j]) <= overlap:\n",
    "            overlap_len += len(current_chunk[j])\n",
    "            j -= 1\n",
    "        # 回退 i 到重叠起始位置\n",
    "        i = start_i + j + 1\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c2a002e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_text = \"\"\"自然语言处理（Natural Language Processing, NLP）是计算机科学和人工智能领域的一个重要研究方向，旨在让计算机能够理解、分析、生成和处理人类语言。NLP结合了语言学、计算机科学和数学等多学科的知识，通过算法和模型，使计算机能够从自然语言文本中提取有用的信息，进行语义理解，并执行相关任务。其应用涵盖了从文本分类、情感分析、机器翻译到自动摘要、信息抽取、对话系统等多个领域。\n",
    "; \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8dd23bca",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[5], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m chunks \u001b[38;5;241m=\u001b[39m \u001b[43msemantic_chunking_with_sliding_window\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample_text\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m分块结果：\u001b[39m\u001b[38;5;124m\"\u001b[39m, chunks)\n",
      "Cell \u001b[1;32mIn[3], line 40\u001b[0m, in \u001b[0;36msemantic_chunking_with_sliding_window\u001b[1;34m(text, chunk_size, overlap)\u001b[0m\n\u001b[0;32m     38\u001b[0m overlap_len \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m     39\u001b[0m j \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(current_chunk) \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m---> 40\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[43mj\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m \u001b[38;5;129;01mand\u001b[39;00m overlap_len \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlen\u001b[39m(current_chunk[j]) \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m overlap:\n\u001b[0;32m     41\u001b[0m     overlap_len \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(current_chunk[j])\n\u001b[0;32m     42\u001b[0m     j \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "chunks = semantic_chunking_with_sliding_window(sample_text)\n",
    "print(\"分块结果：\", chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6145e5ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, chunk in enumerate(chunks):\n",
    "    print(f\"Chunk {i+1} ({len(chunk)} chars):\\n{chunk}\\n{'-'*60}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a31aa60",
   "metadata": {},
   "source": [
    "## 初始化chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856e0f04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 初始化向量数据库\n",
    "client = chromadb.Client()\n",
    "collection = client.get_or_create_collection(\"rag_docs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3da0f62",
   "metadata": {},
   "source": [
    "## 生成嵌入向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "673b1cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedder = SentenceTransformer('all-MiniLM-L6-v2')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e743f7a1",
   "metadata": {},
   "source": [
    "## 存储分块文本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75bf5bf9",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'chunks' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[18], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m embeddings \u001b[38;5;241m=\u001b[39m embedder\u001b[38;5;241m.\u001b[39mencode(\u001b[43mchunks\u001b[49m)\u001b[38;5;241m.\u001b[39mtolist()\n\u001b[0;32m      3\u001b[0m collection\u001b[38;5;241m.\u001b[39madd(\n\u001b[0;32m      4\u001b[0m     embeddings\u001b[38;5;241m=\u001b[39membeddings,\n\u001b[0;32m      5\u001b[0m     documents\u001b[38;5;241m=\u001b[39mchunks,\n\u001b[0;32m      6\u001b[0m     ids\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mid\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(chunks))]\n\u001b[0;32m      7\u001b[0m )\n",
      "\u001b[1;31mNameError\u001b[0m: name 'chunks' is not defined"
     ]
    }
   ],
   "source": [
    "embeddings = embedder.encode(chunks).tolist()\n",
    "\n",
    "collection.add(\n",
    "    embeddings=embeddings,\n",
    "    documents=chunks,\n",
    "    ids=[f\"id{i}\" for i in range(len(chunks))]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f9fd6be",
   "metadata": {},
   "source": [
    "## 检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ffa2e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"语言学\"\n",
    "query_embedding = embedder.encode([query]).tolist()\n",
    "\n",
    "results = collection.query(\n",
    "    query_embeddings=query_embedding,\n",
    "    n_results=2\n",
    ")\n",
    "print(\"Top 2 relevant chunks:\", results['documents'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3221a729",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"自然语言处理\"\n",
    "query_embedding = embedder.encode([query]).tolist()\n",
    "\n",
    "results = collection.query(\n",
    "    query_embeddings=query_embedding,\n",
    "    n_results=2\n",
    ")\n",
    "print(\"Top 2 relevant chunks:\", results['documents'][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "535a930e",
   "metadata": {},
   "source": [
    "## 查询"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e83bf55",
   "metadata": {},
   "source": [
    "- 数据总数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555ddaae",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_docs = collection.count()\n",
    "print(f\"向量数据库中共有 {num_docs} 条数据。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34714f0d",
   "metadata": {},
   "source": [
    "- 数据预览"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a68be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = collection.get(include=[\"documents\", \"embeddings\"])\n",
    "\n",
    "documents = all_data[\"documents\"]\n",
    "embeddings = all_data[\"embeddings\"]\n",
    "ids = all_data[\"ids\"]\n",
    "\n",
    "print(\"\\n全部数据（仅显示前几项预览）:\")\n",
    "for i in range(min(5, len(documents))):  # 控制输出数量避免过长\n",
    "    print(f\"ID: {ids[i]}\")\n",
    "    print(f\"文档内容: {documents[i]}\")\n",
    "    print(f\"嵌入向量（前5维）: {embeddings[:5]}\")   \n",
    "    print(\"-\" * 60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff7b6d2",
   "metadata": {},
   "source": [
    "- 嵌入维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404394cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 嵌入维度（假设所有向量维度相同）\n",
    "try:\n",
    "    embedding_dims = len(embeddings[0])\n",
    "    print(f\"嵌入向量维度：{embedding_dims}\")\n",
    "except Exception as e:\n",
    "    print(f\"无法计算嵌入向量维度，错误信息: {e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
